# -*- coding: utf-8 -*-
import os
import sys
from typing import Collection

import numpy as np
import pandas as pd
import math
import pymongo
import torch
from torch import nn
sys.path.append("/gpfs/scratch/chgwang/XI/Scripts/Refactoring_1/getData")
sys.path.append("/gpfs/scratch/chgwang/XI/Scripts/Refactoring_1/MLModel")

import ParallelNet_1d  # type: ignore
import FedSeq  # type: ignore
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from sklearn.metrics import accuracy_score
from tqdm import tqdm

matplotlib.use("Agg")


def convert_to_one_hot(sour_labels)->list:
    sour_labels = np.array(sour_labels, dtype=np.int32)
    labels_one_hot = np.zeros((6, sour_labels.shape[1]))
    for i in range(sour_labels.shape[1]):
        label_slices = sour_labels[:,i]
        for label in label_slices:
            label = label - 1
            if label != -1:
                labels_one_hot[label, i] = 1
    return labels_one_hot.tolist()

def plot_reponse(id, col):
    data_dict = col.find_one({"_id":id})
    original_labels = convert_to_one_hot(data_dict["label"])
    modeled_labels = data_dict["modeled_label"]
    # aligned length
    original_labels = np.array(original_labels)
    modeled_labels = np.array(modeled_labels)
    assert original_labels.shape[0] == 6
    original_labels = original_labels[:,199:]
    names = ["ia", "ib", "ic"]
    power_data = []
    for name in names:
        power_data.append(np.array(data_dict[name]))
    power_data = np.stack(power_data)
    markers = ["b", "r:", "y--"]
    fig, axes = plt.subplots(nrows=7, ncols=1, sharex=True)
    axes[0].plot(power_data[0,:], markers[0], label="ia")
    axes[0].plot(power_data[1,:], markers[1], label="ib")
    axes[0].plot(power_data[2,:], markers[2], label="ic")
    # axes[0].legend()
    x_labeled = np.arange(199, 601)
    for i in range(6):
        axes[i+1].plot(x_labeled, original_labels[i,:], markers[0], label="experimental result")
        axes[i+1].plot(x_labeled, modeled_labels[i,:], markers[1], label="modeled result")
        # axes[i+1].legend()
    res_fig_dir = "/gpfs/scratch/chgwang/XI/DataBase/Model_figure/Response"
    res_fig_path = os.path.join(res_fig_dir, str(id))
    fig.savefig(res_fig_path)
    plt.close()

# when theta = 0.5, the accuracy_rate = 95.3%

def calc_eval_index(id, col, theta=0.5):
    # should return by channel
    channels_amount = 6
    data_dict = col.find_one({"_id":id})
    if data_dict is None:
        raise Exception
    original_labels = convert_to_one_hot(data_dict["label"])
    modeled_labels = data_dict["modeled_label"]
    # aligned length
    original_labels = np.array(original_labels)
    modeled_labels = np.array(modeled_labels)
    assert original_labels.shape[0] == 6
    original_labels = original_labels[:,199:]
    modeled_labels = np.where(modeled_labels<theta,
                            0, 1)
    modeled_labels = modeled_labels.astype(int)
    original_labels = original_labels.astype(int)
    count_points = modeled_labels.shape[0] * modeled_labels.shape[1]
    eval_index_list = []
    for i in range(channels_amount):
        OP_idx = original_labels[i,:] == 1
        MP_idx = modeled_labels[i,:] == 1
        ON_idx = original_labels[i,:] == 0
        MN_idx = modeled_labels[i,:] == 0

        TP_count = np.count_nonzero(MP_idx & OP_idx)
        FP_count = np.count_nonzero(MP_idx & ON_idx)
        FN_count = np.count_nonzero(MN_idx & OP_idx)
        TN_count = np.count_nonzero(MN_idx & ON_idx)
        eval_index_list.append([TP_count, FP_count, FN_count, TN_count])
    assert count_points == np.sum(eval_index_list)
    return np.array(eval_index_list)

if __name__ == "__main__":
    client = pymongo.MongoClient("mongodb://127.0.0.1:27017/")
    db = client["Power_Fault"]
    col_sour = db["data_sour"]
    # all ID of the data, split data by ID
    ids = col_sour.distinct("_id")
    # for id in ids:
    #     # plot_reponse(id, col_sour)
    theta_list = np.linspace(0, 1, num=61, endpoint=True)
    indices = [(theta, index) for theta in theta_list for index in range(6)]
    indices = pd.MultiIndex.from_tuples(indices, names=["theta", "index"])
    df = pd.DataFrame(index=indices, columns=["Accuracy", "Precision", "Recall", "F1-score"])
    for theta in theta_list:
        eval_array = np.zeros((6,4))
        for id in ids[:]:
            eval_indx_sig = calc_eval_index(id, col_sour, theta)
            eval_array += eval_indx_sig
        for idx, subarr in enumerate(eval_array):
            TP_count = subarr[0]
            FP_count = subarr[1]
            FN_count = subarr[2]
            TN_count = subarr[3]
            accuracy = (TP_count + TN_count) / (TP_count + FP_count + FN_count + TN_count)
            precision = TP_count / (TP_count + FP_count)
            recall = TP_count / (TP_count + FN_count)
            FPR = FP_count / (FP_count + FN_count)
            F1 = 2*precision*recall / (precision + recall)
            df.loc[(theta, idx), "Accuracy"] = accuracy
            df.loc[(theta, idx), "Precision"] = precision
            df.loc[(theta, idx), "Recall"] = recall
            df.loc[(theta, idx), "F1-score"] = F1
            df.loc[(theta, idx), "FPR"] = FPR
    df_dir =  "/gpfs/scratch/chgwang/XI/DataBase/Model_figure"
    df_file_path = os.path.join(df_dir, "evalution_index_theta.csv")
    df.to_csv(df_file_path)
